package com.ruoyi.framework.websocket;

import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * websocket 客户端用户集
 */
@Slf4j
public class WebSocketUsers {
    private static final Map<String, WebSocketSession> USERS = new ConcurrentHashMap<>();

    /**
     * 存储用户
     */
    public static void put(String key, WebSocketSession session) throws IOException {
        WebSocketSession oldSession = USERS.get(key);
        if (oldSession != null) {
            oldSession.close();
        }
        USERS.put(key, session);
    }

    /**
     * 获取在线用户列表
     */
    public static Map<String, WebSocketSession> getUsers() {
        return USERS;
    }

    /**
     * 移除用户
     */
    public static void remove(WebSocketSession session) {
        String key = null;
        boolean flag = USERS.containsValue(session);
        if (flag) {
            Set<Map.Entry<String, WebSocketSession>> entries = USERS.entrySet();
            for (Map.Entry<String, WebSocketSession> entry : entries) {
                WebSocketSession value = entry.getValue();
                if (value.equals(session)) {
                    key = entry.getKey();
                    break;
                }
            }
        }
        remove(key);
    }

    public static void close(String key) throws IOException {
        WebSocketSession session = USERS.get(key);
        if (session != null) {
            session.close();
            remove(key);
        }
    }

    /**
     * 移出用户
     */
    public static void remove(String key) {
        USERS.remove(key);
    }

    /**
     * 群发消息文本消息
     */
    public static void sendMessageToUsersByText(String message) throws IOException {
        Collection<WebSocketSession> values = USERS.values();
        for (WebSocketSession value : values) {
            sendMessageToUserByText(value, message);
        }
    }

    public static void sendMsgToUsersExcludeMeByText(String token, String message) throws IOException {
        Set<String> keys = USERS.keySet();
        for (String key : keys) {
            if (key.equals(token)) {
                continue;
            }
            WebSocketSession value = USERS.get(key);
            sendMessageToUserByText(value, message);
        }
    }

    /**
     * 发送文本消息
     */
    public static void sendMessageToUserByText(WebSocketSession session, String message) throws IOException {
        session.sendMessage(new TextMessage(message));
    }
}
